import pandas as pd
import plotly.express as px
from datetime import datetime
import matplotlib.pyplot as plt
# === 设置真实周期（单位：天） ===
true_period = 0.394

# === 加载 CSV 数据文件 ===
file_path = "/mnt/7b21f1e1-eb25-4cd5-bdb5-06d7d82fa253/Temp/force_photmetry/result50cm/asteroid_photometry_results.csv"
df = pd.read_csv(file_path)

# === 函数：准备指定波段的数据 ===
def prepare_band(df, band):
    df_band = df[(df['band'] == band) & (df['calmag'].notna())].copy()
    df_band['timestamp'] = pd.to_datetime(df_band['obs_time'])
    df_band['mjd'] = (df_band['timestamp'] - df_band['timestamp'].min()).dt.total_seconds() / 86400.0
    df_band['phase'] = (df_band['mjd'] / true_period) % 1
    df_band['tooltip'] = df_band.apply(lambda row: f"time: {row['obs_time']}<br>mag: {row['calmag']:.3f}±{row['magerr']:.3f}", axis=1)
    return df_band

# === 处理 g 和 r 波段 ===
df_r = prepare_band(df, 'r')
df_g = prepare_band(df, 'g')

# 确保两波段在相同时间附近有观测，按最近配对 g 和 r 波段的观测点
from scipy.spatial import cKDTree

# 构建 MJD 和 calmag 映射（仅保留必要列）
r_tbl = df_r[['mjd', 'calmag']].dropna().reset_index(drop=True)
g_tbl = df_g[['mjd', 'calmag']].dropna().reset_index(drop=True)

# 用 KDTree 匹配时间最接近的点
tree_r = cKDTree(r_tbl[['mjd']])
dists, idxs = tree_r.query(g_tbl[['mjd']], distance_upper_bound=0.01)  # 匹配时间差在0.01天内

# 筛选有效匹配点
valid = idxs < len(r_tbl)
matched_g = g_tbl[valid].copy()
matched_r = r_tbl.iloc[idxs[valid]].reset_index(drop=True)

# 计算 g - r 颜色
matched_g['rmag'] = matched_r['calmag']
matched_g['phase'] = (matched_g['mjd'] / true_period) % 1
matched_g['g-r'] = matched_g['calmag'] - matched_g['rmag']

# 按相位绘图
plt.figure(figsize=(7, 4))
plt.scatter(matched_g['phase'], matched_g['g-r'], color='purple', alpha=0.7)
plt.xlabel("Phase (folded by 0.394 days)")
plt.ylabel("g - r Color")
plt.title("Color Variation with Phase")
plt.grid(True)
plt.tight_layout()
plt.savefig("color-color.png",dpi=200)
